import torch
import cv2
import os
import sys
import numpy as np
import time
import gc
#from queue import Queue
from multiprocessing import Process, Queue
from ultralytics import YOLO
from threading import Thread
import logging
from nodes.Node import Node
from ultralytics.utils.plotting import Annotator, colors

class FishFilterNode(Node, Process):
    def __init__(self,name="FishFilterNode",parent=None):
        super().__init__(name, parent)
        super(Process,self).__init__()

        self.is_stop=False

    def run(self):
        self.setStatus(Node.NODE_STATUS_RUNNING)
        self.is_stopped = False

        nextFishId = 0
        fishMap = {}
        objLife = {}

        while True:
            if self.is_stopped:
                break
            else:
                if( not self._framePool.empty()):

                    frame, videoIndex, frameIndex, tracked_dets = self._framePool.get()

                    #annotator = Annotator(frame, line_width=4)

                    if len(tracked_dets)>0:
                        bbox_xyxy = tracked_dets[:,:4]
                        identities = tracked_dets[:, 8]
                        categories = tracked_dets[:, 4]

                        for i, box in enumerate(bbox_xyxy):
                            clsIdx = int(categories[i])
                            objIdx = int(identities[i])

                            fishId = 0

                            if objIdx not in fishMap:
                                fishMap[objIdx]=nextFishId
                                fishId = nextFishId
                                nextFishId += 1
                            else:
                                fishId = fishMap[objIdx]

                            tracked_dets[i, 8]=fishId

                            objLife[objIdx]=4


                            #if( categories[i] == 0 ):
                            # annotator.box_label(box.tolist(), str(fishId), color=colors(0, True))
                            #else:
                            #annotator.box_label(box.tolist(),labels[clsIdx] + ":" + str(objIdx),color=colors(clsIdx,True))

                    for oid in objLife:
                        objLife[oid]-=1
                        if objLife[oid] < 0:
                            del objLife[oid]
                            del fishMap[oid]
                            break

                    #print("Len objLife:" + str(len(objLife)))

                    if frame is not None:
                        self.passFrame([frame, videoIndex, frameIndex, tracked_dets])
                        logging.info("FishFilter filter a frame.")


            if self._framePool.empty():
                gc.collect()
                if self.getParentStatus() == Node.NODE_STATUS_EXITED:
                    logging.error("Parent of FishFilterNode is completed. FishFilterNode is going to close.")
                    break
                #else:
                #    logging.info("Fish Filter Pool is empty......")
            #else:
            #    logging.info("Fish Filter Pool is not empty.")

        self.setStatus(Node.NODE_STATUS_EXITED)
        logging.error("=========>>>>>>>>>>>>>>>>>>>>")
        logging.error("FishFilterNode is completed.")
        #sys.exit()

    def stop (self):
        self.is_stop = True
        self._parent.stop()
